import numpy as np
import math
import pygad
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False
plt.switch_backend('TkAgg')

# -------------------------- 1. 基础参数与复用函数 --------------------------
target_true = np.array([0, 200, 0])  # 真目标
r, h = 7, 10  # 真目标的半径7m、高10m
target_fake = np.array([0, 0, 0])  # 假目标
missile_M1 = np.array([20000, 0, 2000])  # M1初始位置
missile_speed = 300  # 导弹速度
drone_FY1 = np.array([17800, 0, 1800])  # FY1初始位置
g = 9.8  # 重力加速度
smoke_radius = 10  # 烟幕半径
smoke_sink_speed = 3  # 下沉速度
smoke_valid_time = 20  # 有效时间
drone_speed_range = [70, 140]  # 无人机速度约束

# 计算导弹飞到假目标的时间（关键约束）
missile_distance_to_fake = np.linalg.norm(target_fake - missile_M1)
missile_time_to_fake = missile_distance_to_fake / missile_speed  # 导弹到达假目标的时间


# 定义获取长方体顶点的函数
def get_cylinder_bounding_box_vertices(target_center, radius, height):
    x_min = target_center[0] - radius
    x_max = target_center[0] + radius
    y_min = target_center[1] - radius
    y_max = target_center[1] + radius
    z_min = target_center[2]
    z_max = target_center[2] + height
    #长方体顶点
    vertices = [
        np.array([x_min, y_min, z_min]),
        np.array([x_max, y_min, z_min]),
        np.array([x_max, y_max, z_min]),
        np.array([x_min, y_max, z_min]),
        np.array([x_min, y_min, z_max]),
        np.array([x_max, y_min, z_max]),
        np.array([x_max, y_max, z_max]),
        np.array([x_min, y_max, z_max])
    ]
    return vertices


# 预计算目标顶点
target_vertices = get_cylinder_bounding_box_vertices(target_true, r, h)


def missile_real_time_position(missile_init_pos, t):
    # 导弹到达假目标后不再移动
    if t >= missile_time_to_fake:
        return target_fake
    if t < 0:
        return missile_init_pos
    dir_vec = target_fake - missile_init_pos
    dir_norm = np.linalg.norm(dir_vec)
    if dir_norm <= 1e-6:
        return missile_init_pos
    dir_unit_vec = dir_vec / dir_norm
    return missile_init_pos + missile_speed * t * dir_unit_vec


# -------------------------- 2. 烟幕实时位置计算 --------------------------
def get_smoke_explosion_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate):
    # 投放点坐标
    drop_x = drone_init_pos[0] + drone_speed * math.cos(theta) * t_drop
    drop_y = drone_init_pos[1] + drone_speed * math.sin(theta) * t_drop
    drop_z = drone_init_pos[2]
    # 烟幕弹平抛运动
    explosion_x = drop_x + drone_speed * math.cos(theta) * t_detonate
    explosion_y = drop_y + drone_speed * math.sin(theta) * t_detonate
    # 竖直自由下落
    fall_distance = 0.5 * g * (t_detonate ** 2)
    explosion_z = max(drop_z - fall_distance, 0.0)
    return np.array([explosion_x, explosion_y, explosion_z])


def smoke_real_time_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate, t):
    detonate_t = t_drop + t_detonate
    if t < detonate_t - 1e-6:
        return np.array([0, 0, 0])
    explosion_center = get_smoke_explosion_center(drone_init_pos, drone_speed, theta, t_drop, t_detonate)
    sink_distance = smoke_sink_speed * (t - detonate_t)
    real_z = max(explosion_center[2] - sink_distance, 0.0)
    return np.array([explosion_center[0], explosion_center[1], real_z])


# -------------------------- 3. 遮蔽判断 --------------------------
def is_segment_sphere_intersect(missile_pos, vertex_pos, smoke_center, smoke_radius):
    vec_AC = smoke_center - missile_pos
    vec_AB = vertex_pos - missile_pos
    ab_sq_mag = np.dot(vec_AB, vec_AB)
    if ab_sq_mag < 1e-12:
        return False
    k = np.dot(vec_AC, vec_AB) / ab_sq_mag
    if k <= 0.0:
        min_dist = np.linalg.norm(vec_AC)
    elif 0.0 < k < 1.0:
        cross_prod = np.cross(vec_AC, vec_AB)
        min_dist = np.linalg.norm(cross_prod) / np.sqrt(ab_sq_mag)
    else:
        vec_BC = smoke_center - vertex_pos
        min_dist = np.linalg.norm(vec_BC)
    return min_dist <= smoke_radius + 1e-6


def is_effective_occlusion(t, drone_speed, theta, t_drop, t_detonate):
    detonate_t = t_drop + t_detonate
    smoke_start = detonate_t
    smoke_end = detonate_t + smoke_valid_time
    if t < smoke_start - 1e-6 or t > smoke_end + 1e-6:
        return False
    missile_pos = missile_real_time_position(missile_M1, t)
    smoke_center = smoke_real_time_center(drone_FY1, drone_speed, theta, t_drop, t_detonate, t)
    for vertex in target_vertices:
        if is_segment_sphere_intersect(missile_pos, vertex, smoke_center, smoke_radius):
            return True
    return False


# -------------------------- 4. 有效遮蔽时长计算 --------------------------
def calculate_effective_duration(drone_speed, theta, t_drop, t_detonate, time_step=0.01):
    detonate_t = t_drop + t_detonate
    smoke_start = detonate_t
    smoke_end = detonate_t + smoke_valid_time
    total_duration = 0.0
    prev_effective = False
    t = smoke_start
    while t <= smoke_end + 1e-6:
        current_effective = is_effective_occlusion(t, drone_speed, theta, t_drop, t_detonate)
        if current_effective and not prev_effective:
            start_t = t
        elif not current_effective and prev_effective:
            total_duration += t - start_t
        prev_effective = current_effective
        t += time_step
    if prev_effective:
        total_duration += smoke_end - start_t
    return round(total_duration, 3)


# -------------------------- 5. 适应度函数（核心优化逻辑） --------------------------
def fitness_func(ga_instance, solution, solution_idx):
    drone_speed, theta, t_drop, t_detonate = solution
    # 参数约束
    if not (drone_speed_range[0] <= drone_speed <= drone_speed_range[1]):
        return 0
    detonate_t = t_drop + t_detonate
    if detonate_t > missile_time_to_fake:
        return 0
    duration = calculate_effective_duration(drone_speed, theta, t_drop, t_detonate)
    return duration


# -------------------------- 6. 用遗传算法训练 --------------------------
# 用于保存每代最大遮蔽时长
generation_best_fitness = []

def on_generation(ga_instance):
    # 获取当前种群最大适应度
    best_fitness = np.max(ga_instance.last_generation_fitness)
    generation_best_fitness.append(best_fitness)
    print(f"当前迭代次数：{ga_instance.generations_completed}, 最大遮蔽时长：{best_fitness:.2f}s")


if __name__ == "__main__":
    gene_space = [
        {'low': 70, 'high': 140},  # 速度
        {'low': 0, 'high': 2 * math.pi},  # 方向角
        {'low': 0, 'high': min(30, missile_time_to_fake)},  # 投放延迟不超过导弹到达假目标的时间
        {'low': 0, 'high': min(10, missile_time_to_fake)}  # 起爆延迟不超过导弹到达假目标的时间
    ]

    ga_instance = pygad.GA(
        num_generations=100,  # 迭代次数
        num_parents_mating=10,  # 每代参与交配的父代数量
        fitness_func=fitness_func,  # 适应度函数
        sol_per_pop=50,  # 种群大小
        num_genes=4,  # 变量数量（4个优化参数）
        gene_space=gene_space,  # 变量范围
        parent_selection_type="sss",  # 选择策略：稳态选择
        crossover_type="single_point",  # 交叉策略：单点交叉
        mutation_type="random",  # 变异策略：随机变异
        mutation_percent_genes=5,  # 变异基因比例（5%）
        mutation_num_genes=1,  # 直接指定变异1个基因
        parallel_processing=8,  # 启用并行计算
        on_generation=on_generation  # 每一代的回调函数
    )

    print("开始遗传算法优化...")
    ga_instance.run()

    # 结果解析
    solution, solution_fitness, solution_idx = ga_instance.best_solution()
    drone_speed_opt, theta_opt, t_drop_opt, t_detonate_opt = solution

    dir_x_opt = math.cos(theta_opt)
    dir_y_opt = math.sin(theta_opt)
    detonate_time_opt = t_drop_opt + t_detonate_opt

    # 计算投放点和起爆点
    drop_x = drone_FY1[0] + drone_speed_opt * math.cos(theta_opt) * t_drop_opt
    drop_y = drone_FY1[1] + drone_speed_opt * math.sin(theta_opt) * t_drop_opt
    drop_z = drone_FY1[2]
    drop_point = (drop_x, drop_y, drop_z)

    explosion_point = get_smoke_explosion_center(
        drone_FY1, drone_speed_opt, theta_opt, t_drop_opt, t_detonate_opt
    )

    # 方向角度（0~360度，逆时针为正，x轴为正向）
    theta_deg = (theta_opt * 180 / math.pi) % 360

    print("\n最优解：")
    print(f"1. 飞行参数：速度={drone_speed_opt:.2f}m/s")
    print(f"2.方向角度={theta_deg:.2f}°（以x轴为正向，逆时针方向为正，0~360度）")
    print(f"3.投放参数：投放延迟={t_drop_opt:.2f}s，起爆延迟={t_detonate_opt:.2f}s")
    print(f"  起爆时间：{detonate_time_opt:.2f}s")
    print(f"4.烟幕干扰弹投放点=({drop_point[0]:.2f}, {drop_point[1]:.2f}, {drop_point[2]:.2f})")
    print(f"5.烟幕干扰弹起爆点=({explosion_point[0]:.2f}, {explosion_point[1]:.2f}, {explosion_point[2]:.2f})")
    print(f"6. 最大遮蔽时长：{solution_fitness:.2f}s")

    # 绘制折线图
    plt.figure()
    plt.plot(range(1, len(generation_best_fitness)+1), generation_best_fitness, marker='o')
    plt.xlabel("迭代次数")
    plt.ylabel("最大遮蔽时长 (s)")
    plt.title("遗传算法迭代过程最大遮蔽时长变化")
    plt.grid(True)
    plt.show()